#!/beegfs/group/lqcd/software/anaconda3/bin/python3
# encoding: utf-8
import numpy as np
import xml.etree.ElementTree as ET
import os
import time


def ReadWlpXML(filename):
    if type(filename) is list:
        nconf = len(filename)
        # read Params
        tree = ET.parse(filename[0])
        root = tree.getroot()

        length = {}
        for i in root.findall("./InlineObservables/elem/WilsonLoop/wils_loop3/"):
            if "length" in i.tag:
                length[i.tag] = int(i.text)
        WilsonLoop3 = np.zeros((nconf, length["lengtht"], length["length"]))

        for i in root.findall("./InlineObservables/elem/WilsonLoop/wils_loop2/"):
            if "length" in i.tag:
                length[i.tag] = int(i.text)
        WilsonLoop2 = np.zeros((nconf, length["lengtht"], length["lengthr"]))

        for i in root.findall("./InlineObservables/elem/WilsonLoop/wils_loop1/"):
            if "length" in i.tag:
                length[i.tag] = int(i.text)
        WilsonLoop1 = np.zeros((nconf, length["lengthr"], length["lengthr"]))

        for iconf in range(nconf):
            tree = ET.parse(filename[iconf])
            root = tree.getroot()

            ir = 0
            for iloop in root.findall("./InlineObservables/elem/WilsonLoop/wils_loop1/wloop1/elem/loop"):
                WilsonLoop1[iconf, :, ir] = np.fromstring(iloop.text, sep=' ')
                ir = ir+1

            ir = 0
            for iloop in root.findall("./InlineObservables/elem/WilsonLoop/wils_loop2/wloop2/elem/loop"):
                WilsonLoop2[iconf, :, ir] = np.fromstring(iloop.text, sep=' ')
                ir = ir+1

            ir = 0
            for iloop in root.findall("./InlineObservables/elem/WilsonLoop/wils_loop3/wloop3/elem/loop"):
                WilsonLoop3[iconf, :, ir] = np.fromstring(iloop.text, sep=' ')
                ir = ir+1
        return WilsonLoop1, WilsonLoop2, WilsonLoop3


def main():
    smear = ""  # mean do HYP_smearing

    conf = np.arange(1, 201)
    betalist = ['purgauge_2432']

    # for beta in ['beta24','beta28','beta30']:
    for beta in betalist:
        path = "./purgauge_2432/output/"
        filename = os.listdir("./purgauge_2432/output/")
        filename = [path + i for i in filename]
        filename.sort()
        startT = time.time()
        nconf = len(filename)
        WilsonLoop1, WilsonLoop2, WilsonLoop3 = ReadWlpXML(filename)
        np.savez(beta+"_WilsonLoop%s" %
                 smear, WilsonLoop1, WilsonLoop2, WilsonLoop3)
        with open("wilsonloop1.txt", "ab") as f:
            for iconf in range(nconf):
                np.savetxt(f, np.array(["# conf %i" % iconf]), fmt="%s")
                np.savetxt(f, np.array(["#space-like wilson loop"]), fmt="%s")
                np.savetxt(f, WilsonLoop1[iconf, :, :])
        with open("wilsonloop2.txt", "ab") as f:
            for iconf in range(nconf):
                np.savetxt(f, np.array(["# conf %i" % iconf]), fmt="%s")
                np.savetxt(f, np.array(["#time-like wilson loop"]), fmt="%s")
                np.savetxt(f, WilsonLoop2[iconf, :, :])
        with open("wilsonloop3.txt", "ab") as f:
            for iconf in range(nconf):
                np.savetxt(f, np.array(["# conf %i" % iconf]), fmt="%s")
                np.savetxt(f, np.array(
                    ["#non-axias time-like wilson loop"]), fmt="%s")
                np.savetxt(f, WilsonLoop3[iconf, :, :])
        print("all time used is ", time.time()-startT)


if __name__ == '__main__':
    main()
